# -*- coding: UTF-8 -*-
"""
@Date    ：2025/9/10 15:09 
@Author  ：Liu Yuezhao
@Project ：bert 
@File    ：utils.py
@IDE     ：PyCharm 
"""
import pandas as pd
import yaml
import json
import numpy as np
import random, os
import torch
from torch import nn

def load_config(file_path):
    """
    读取yaml配置文件
    :param file_path: 配置文件路径
    :return: 配置信息
    """
    with open(file_path, 'r', encoding='utf-8') as file:
        try:
            config_data = yaml.safe_load(file)  # 使用 safe_load 避免安全风险
            return config_data
        except yaml.YAMLError as exc:
            print(f"YAML解析错误: {exc}")
            return None

def read_json(path, encoding='utf-8'):
    """
    读取json文件
    Args:
        path (_type_): 文件路径.
        encoding (str, optional): 编码方式.
    Returns:
        data: 文件内容
    """
    print(f'start loading json file: {path}')

    with open(path, 'r', encoding=encoding) as f:
        data = json.load(f)

    return data

def write_json(data, path):
    """
    写入json文件
    :param data: 写入数据
    :param path: 文件路径
    """
    with open(path, encoding='utf-8', mode='w+') as f:
        json.dump(data, f, indent=4)

    print(f'finish dump json str to path: {path}')

def split_data_to_train_valid_test(data: pd.DataFrame, train_valid_ratio: list = [0.7, 0.2]):
    assert sum(train_valid_ratio) < 1, "You must left test samples"
    print("Start to split dataset")
    data = data.sample(frac=1).reset_index(drop=True)
    data_len = data.shape[0]
    train_size = int(train_valid_ratio[0] * data_len)
    valid_size = int(train_valid_ratio[1] * data_len)
    train_data = data.iloc[:train_size, :].reset_index(drop=True)
    valid_data = data.iloc[train_size:train_size + valid_size, :].reset_index(drop=True)
    test_data = data.iloc[train_size + valid_size:, :].reset_index(drop=True)
    return train_data, valid_data, test_data

def set_seed(seed):
    """
    设置所有相关随机库的种子，以确保实验的可复现性。
    Args:
        seed (int): 随机种子值。默认为 42。
    """
    # 1. Python 内置 random 模块
    random.seed(seed)

    # 2. NumPy
    np.random.seed(seed)

    # 3. PyTorch - CPU
    torch.manual_seed(seed)

    # 4. PyTorch - GPU (所有可见的 CUDA 设备)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)  # 多 GPU 的情况下

    # 5. 设置 cuDNN 确定性模式（可能会影响性能）
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # 6. 设置 PyTorch 使用确定性算法（PyTorch >= 1.7）
    # 注意：这可能限制某些高性能算法的使用
    try:
        torch.use_deterministic_algorithms(True)
    except:
        print("警告：当前 PyTorch 版本不支持 `torch.use_deterministic_algorithms`。")

    # 7. 设置 cuBLAS 的确定性工作区（需要设置环境变量）
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'  # 或 ':16:8'

    # 8. 禁止 TensorFlow/其他库的干扰（可选）
    os.environ['PYTHONHASHSEED'] = str(seed)

    print(f"随机种子已设置为 {seed}，所有随机性操作将被控制。")

def save_model(model: nn.Module, epoch: int, save_path: str):
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
        },
        f"{save_path}_{epoch}.pt",
    )
    print(f"save model to path {save_path} successfully !")

# 正负样本采样
def undersample_negative_class(df, target_col='target_1', pos_label=1, neg_label=0, ratio=4):
    """
    按照 pos:neg = 1:ratio 的比例对负样本进行欠采样
    """
    # 分离正负样本
    pos_df = df[df[target_col] == pos_label]
    neg_df = df[df[target_col] == neg_label]

    print(f"1类别样本：{len(pos_df)}=====0类别样本：{len(neg_df)}")

    # 计算所需负样本数量
    num_neg_samples = len(pos_df) * ratio

    # 如果负样本不够，就全部保留
    if num_neg_samples >= len(neg_df):
        undersampled_neg_df = neg_df
    else:
        undersampled_neg_df = neg_df.sample(n=num_neg_samples, random_state=42)

    print("------采样后------")
    print(f"1类别样本：{len(pos_df)}=====0类别样本：{len(undersampled_neg_df)}")
    # 合并并打乱顺序
    undersampled_df = pd.concat([pos_df, undersampled_neg_df], axis=0).sample(frac=1, random_state=42).reset_index(drop=True)

    return undersampled_df